"""
Reinforcement learning maze example.
Red rectangle:          explorer.
Black rectangles:       hells       [reward = -1].
Yellow bin circle:      paradise    [reward = +1].
All other states:       ground      [reward = 0].
This script is the environment part of this example. The RL is in RL_brain.py.
View more on my tutorial page: https://morvanzhou.github.io/tutorials/
"""

import numpy as np
import time
import tkinter as tk
from PIL import Image, ImageTk

UNIT = 40  # pixels
MAZE_H = 5  # grid height
MAZE_W = 5  # grid width
image = None
photo = None
sleep_time = 0.00001


class Maze(tk.Tk, object):
    def __init__(self):
        """初始化参数"""
        super(Maze, self).__init__()
        self.action_space = []
        self.action_d = ['u', 'd', 'l', 'r']
        for i in range(len(self.action_d)):
            for j in range(len(self.action_d)):
                self.action_space.append([self.action_d[i] + self.action_d[j]])
        self.n_actions = len(self.action_space)
        self.title('maze')
        self.robot = []
        self.barrier_center = []
        self.treasure = []
        self.treasure_center = []
        self.terminus_center = None
        self.geometry('{0}x{1}'.format(MAZE_H * UNIT, MAZE_H * UNIT))
        self._build_maze()

    def creat_img(self, origin, abscissa, ordinate, file):
        center = origin + np.array([UNIT * abscissa, UNIT * ordinate])
        return self.canvas.create_image(
            center[0], center[1], image=file)

    def creat_rec(self, origin, abscissa, ordinate, color):
        center = origin + np.array([UNIT * abscissa, UNIT * ordinate])
        if color == 'black':
            self.barrier_center.append([center[0], center[1]])
        elif color == 'yellow':
            self.treasure_center.append([center[0], center[1]])
        elif color == 'green':
            self.terminus_center = [center[0], center[1]]
        return self.canvas.create_rectangle(
            center[0] - 15, center[1] - 15,
            center[0] + 15, center[1] + 15,
            fill=color)

    def _build_maze(self):
        """构建迷宫"""
        self.canvas = tk.Canvas(self, bg='white',
                                height=MAZE_H * UNIT,
                                width=MAZE_W * UNIT)
        # create grids
        for c in range(0, MAZE_W * UNIT, UNIT):
            x0, y0, x1, y1 = c, 0, c, MAZE_H * UNIT
            self.canvas.create_line(x0, y0, x1, y1)
        for r in range(0, MAZE_H * UNIT, UNIT):
            x0, y0, x1, y1 = 0, r, MAZE_W * UNIT, r
            self.canvas.create_line(x0, y0, x1, y1)

        # create origin
        origin = np.array([20, 20])

        # barrier
        self.barrier1 = self.creat_rec(origin, 4, 0, 'black')
        self.barrier2 = self.creat_rec(origin, 1, 1, 'black')
        self.barrier3 = self.creat_rec(origin, 2, 2, 'black')
        self.barrier4 = self.creat_rec(origin, 4, 2, 'black')
        self.barrier5 = self.creat_rec(origin, 0, 4, 'black')
        self.barrier6 = self.creat_rec(origin, 3, 4, 'black')

        # create treasure
        self.treasure.append(self.creat_rec(origin, 1, 4, 'yellow'))
        self.treasure.append(self.creat_rec(origin, 2, 1, 'yellow'))

        # create robot
        global image
        global photo
        image = Image.open("robot.png")
        photo = ImageTk.PhotoImage(image)
        self.robot.append(self.creat_img(origin, 0, 1, photo))
        self.robot.append(self.creat_img(origin, 1, 0, photo))

        # create red rect
        self.rect = self.canvas.create_rectangle(
            origin[0] - 15, origin[1] - 15,
            origin[0] + 15, origin[1] + 15,
            fill='red')

        # create terminus
        self.terminus = self.creat_rec(origin, 4, 4, 'green')
        # pack all
        self.canvas.pack()

    def reset(self):
        """清空图像，还原最初状态"""
        self.update()
        time.sleep(sleep_time)

        # 删掉机器人和宝藏
        for i in range(len(self.robot)):
            self.canvas.delete(self.robot[i])
            self.canvas.delete(self.treasure[i])

        # 重新画宝藏
        origin = np.array([20, 20])
        self.robot = []
        self.robot.append(self.creat_img(origin, 0, 1, photo))
        self.robot.append(self.creat_img(origin, 1, 0, photo))
        self.treasure = []
        self.treasure_center = []
        self.treasure.append(self.creat_rec(origin, 1, 4, 'yellow'))
        self.treasure.append(self.creat_rec(origin, 2, 1, 'yellow'))

        # return observation 机器人两个的坐标
        observation = [self.canvas.coords(robot) for robot in self.robot]

        # 添加找到宝藏状态
        # 0 没有拿到
        # 1 拿到第一个宝藏
        # 2 拿到第二个宝藏
        # 3 拿到第三个宝藏
        observation.append(0)
        return observation

    def step(self, s, action):
        base_action = np.array([[0, 0]] * len(s))
        action = [int(action / 4), action % 4]
        for i in range(2):
            if s[i] == 'Terminal':
                continue
            if action[i] == 0:  # up
                if s[i][1] > UNIT:
                    base_action[i][1] -= UNIT
            elif action[i] == 1:  # down
                if s[i][1] < (MAZE_H - 1) * UNIT:
                    base_action[i][1] += UNIT
            elif action[i] == 2:  # left
                if s[i][0] > UNIT:
                    base_action[i][0] -= UNIT
            elif action[i] == 3:  # right
                if s[i][0] < (MAZE_W - 1) * UNIT:
                    base_action[i][0] += UNIT

        # move agent
        for i in range(2):
            self.canvas.move(self.robot[i], base_action[i][0], base_action[i][1])

        s_ = [self.canvas.coords(robot) for robot in self.robot]  # next state

        # reward function

        # 碰撞终止，-100分
        if s_[0] == s_[1] and s_[0] != self.terminus_center:
            s_.append(s[2])
            return s_, -20, True
        done = False
        reward = 0
        # 初始宝藏状态
        treasure_sta = s[2]

        for i in range(len(s_)):
            if s[i] == 'Terminal':
                s_[i] = s[i]
                continue
            if s_[i] in self.treasure_center:
                for j in range(len(self.treasure_center)):
                    if ((treasure_sta & (1 << j)) == 0) and s_[i] == self.treasure_center[j]:
                        self.update()
                        self.canvas.delete(self.treasure[j])
                        treasure_sta |= 1 << j
                        reward += 30
            elif s_[i] in self.barrier_center:
                reward += -10
                done = True
            elif s_[i] == self.terminus_center:
                s_[i] = 'Terminal'
                reward += 60
            else:
                reward += -1
        if s_[0] == s_[1]:
            done = True

        # 宝藏状态放进去
        s_.append(treasure_sta)

        return s_, reward, done

    def render(self):
        time.sleep(sleep_time)
        self.update()


def update():
    for t in range(10):
        env.reset()
        while True:
            env.render()
            a = 7
            s, r, done = env.step(a)
            if done:
                break


if __name__ == '__main__':
    env = Maze()
    # env.reset()
    env.after(100, update)
    env.mainloop()
